import os

import torch
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers.utils import load_image

from util import seed_everything

os.environ["http_proxy"] = "http://192.168.3.116:7890/"
os.environ["https_proxy"] = "http://192.168.3.116:7890/"

image_path="data/image/3.png"
image_path="data/12d/00000.png"
init_image = load_image("data/image/3.png")

seed_everything(1)

model_key = "runwayml/stable-diffusion-v1-5"
# pipe = StableDiffusionPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_key, torch_dtype=torch.float16).to("cuda")

lora_safetensors_path = "lora/1.5/add_detail.safetensors"
pipe.load_lora_weights(lora_safetensors_path)
pipe.fuse_lora(lora_scale=0.9)

prompt = "godrays,In an interesting hyper detailed masterpiece, dynamic realistic digital art, awesome quality,elevation market,gravity celestial parallax,evergreen curiouser and curiouser caustics rendering,henon map, order,inorganic,oxygen-rich air.,(high quality:1.3),(best quality:1.3),(masterpiece:1.3),official wallpaper,4k textures, epic(1.2),(extremely detailed:1.01),(sharp focus:1.01),(hdr:1.01),,"
prompt = "godrays,In an interesting hyper detailed masterpiece, dynamic realistic digital art, awesome quality,elevation market,gravity celestial parallax,evergreen curiouser and curiouser caustics rendering,henon map, order,inorganic,oxygen-rich air.,(high quality:1.3),(best quality:1.3),(masterpiece:1.3),official wallpaper,4k textures, epic(1.2),(extremely detailed:1.01),(sharp focus:1.01),(hdr:1.01),"
negative_prompt = "(worst quality:2), (bad quality:2), (normal quality:2), lowers, bad anatomy, bad hands, (multiple views), human"

init_images = "data/12f"
if os.path.exists(init_images):
    init_image_list = [name for name in os.listdir(init_images) if os.path.isfile(os.path.join(init_images, name))]
    for init_image in init_image_list:
        image = pipe(prompt, image=load_image(os.path.join(init_images, init_image)), guidance_scale=7.5).images[0]
        image.save(os.path.join("image","test1",init_image))



